from asyncore import read
from copy import copy
from datetime import datetime
from io import StringIO
import json
import os
import sys
import time
import psycopg2

from data_reader.DataReader import DataReader
from data_reader.NsReader import NsReader
from data_reader.OrReader import OrReader
from utils.Timer import Timer

class PgFlat:
    def drop_table(self, conn):
        cur = conn.cursor()
        cur.execute('DROP TABLE point_cloud')
        conn.commit()
        cur.close()

    def create_table(self, conn):
        cur = conn.cursor()
        cur.execute("CREATE TABLE point_cloud( "
                    "x    REAL  NOT NULL,"
                    "y    REAL  NOT NULL,"
                    "z    REAL  NOT NULL,"
                    "sweep INTEGER NOT NULL,"
                    "tstamp TIMESTAMP NOT NULL)")
        conn.commit()
        cur.close()

    def create_index(self, conn):
        cur = conn.cursor()
        cur.execute("CREATE INDEX sweep_index ON point_cloud(sweep)")
        cur.execute("CREATE INDEX tstamp_index ON point_cloud(tstamp)")
        conn.commit()
        cur.close()

    def load_data(self, test_case):
        timer = Timer()
        timer.Start()
        reader_name = test_case["reader"]
        data_dir = test_case["dir"]
            
        loadermodule = "data_reader." + reader_name
        exec('from ' + loadermodule + ' import ' + reader_name)
        reader_class = getattr(sys.modules[loadermodule], reader_name)
        reader = reader_class(data_dir)
            
        conn = psycopg2.connect("dbname=postgres user=postgres password=postgres")
        self.drop_table(conn)
        self.create_table(conn)
        cur = conn.cursor()
        timer.StopAndPrint("Init")

        sweep_index = 0
        while reader.has_more_sweep():
            sweep_index += 1
            points = reader.read_sweep()
            tstamp = reader.read_timestamp()
            timer.Start()
            copy_str = ''
            for p in points:
                copy_str += "%f\t%f\t%f\t%d\t%s\n"%(p[0], p[1], p[2], sweep_index, datetime.fromtimestamp(tstamp))

            # Execute copy
            cur.copy_from(StringIO(copy_str), 'point_cloud', columns=('x', 'y', 'z', 'sweep', 'tstamp'))
            conn.commit()
            timer.Pause()
        timer.StopAndPrint("Load")
        timer.Start()
        self.create_index(conn)
        cur.close()
        conn.close()
        timer.StopAndPrint("Close")

if __name__ == '__main__':
    with open(os.getcwd() + '/config.json', 'r') as f:
        config = json.load(f)

    data_base_path = config["data_base_path"]
    test_case = config["test_case"]
    for test_case in config["test_case"]:
        test_case = config["collect_dataset"][test_case]
        PgFlat().load_data(test_case)
        